#include <vector>
#include <iostream>
#include <thread>
#include <cmath>

// input must be a vector
void softmax_kernel(float *input, unsigned int size, float *output) {
    float sum = 0.0f;
    for (int i = 0; i < size; i++) {
        sum += exp(input[i]);
    }

    for (int i = 0; i < size; i++) {
        output[i] = exp(input[i]) / sum;
    }
    
}

struct SoftmaxArgs {
    int start;
    int end;
    float *input;
    unsigned int size;
    float *output;
    float *temp;
}*s;

void softmax_thread_function(SoftmaxArgs *s) {
    for (int i = s->start; i < s->end; i++) {
        int sw = i * s->size; 
        float *inptr = &s->input[sw];
        float *outptr = &s->output[sw];
        softmax_kernel(inptr, s->size, outptr);
    }
}

void softmax(float *input, unsigned int x_size, unsigned int y_size, float *output) {
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<SoftmaxArgs> mes(thread_num);
    int blocksdim;
    if (x_size%thread_num != 0) {
        if (x_size < thread_num) {
            for (int i = 0; i < x_size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = SoftmaxArgs{start, end, input, y_size, output, NULL};
            }
        }else{
            blocksdim = x_size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = x_size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = SoftmaxArgs{start, end, input, y_size, output, NULL};
            }
        }
    }else{
        blocksdim = x_size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = SoftmaxArgs{start, end, input, y_size, output, NULL};
            }
    }
    if (x_size < thread_num) {
        thread_num = x_size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(softmax_thread_function, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}


// backward
// input must be a vector
void softmax_back_kernel(float *input, unsigned int size, float *output, float *loss) {
    float sum = 0.0f;
    for (int i = 0; i < size; i++) {
        sum += exp(input[i]);
    }
    for (int i = 0; i < size; i++) {
        float si = exp(input[i]);
        for (int j = 0; j < size; j++) {
            float sj = exp(input[j]);
            if (i == j) {
                output[i] += si - si*si;
            }
            else {
                output[i] += -si*sj;
            }
        }
    }

    for (int i = 0; i < size; i++) {
        output[i] *= loss[i];
    }
}

void softmax_back_thread_function(SoftmaxArgs *s) {
    for (int i = s->start; i < s->end; i++) {
        int sw = i * s->size; 
        float *inptr = &s->input[sw];
        float *outptr = &s->output[sw];
        float *lossptr = &s->temp[sw];
        softmax_back_kernel(inptr, s->size, outptr, lossptr);
    }
}

void softmax_backward(float *input, unsigned int x_size, unsigned int y_size, float *output, float *loss){
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<SoftmaxArgs> mes(thread_num);
    int blocksdim;
    if (x_size%thread_num != 0) {
        if (x_size < thread_num) {
            for (int i = 0; i < x_size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = SoftmaxArgs{start, end, input, y_size, output, loss};
            }
        }else{
            blocksdim = x_size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = x_size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = SoftmaxArgs{start, end, input, y_size, output, loss};
            }
        }
    }else{
        blocksdim = x_size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = SoftmaxArgs{start, end, input, y_size, output, loss};
            }
    }
    if (x_size < thread_num) {
        thread_num = x_size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(softmax_back_thread_function, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}

int main() {

    float input[5] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
    float output[5] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    softmax_kernel(input, 5, output);

    for (int i = 0; i < 5; i++) {
        std::cout << output[i] << std::endl;
    }
    return 0;
}